import sys
import torch
import numpy as np
import cv2
from pathlib import Path

# === 1️⃣ 引入 UltraSam 源码路径 ===
ULTRASAM_PATH = Path("/data/suziren25/suziren/ultra-sam-light/models")
if str(ULTRASAM_PATH) not in sys.path:
    sys.path.append(str(ULTRASAM_PATH))

# === 2️⃣ 导入 UltraSam 模型 ===
# 按 UltraSam 官方仓库结构导入主模型
try:
    from ultrasam import UltraSAM  # UltraSam-main/models/ultrasam.py
except ModuleNotFoundError:
    raise ImportError(
        f"未找到 UltraSAM 模型，请确认路径正确: {ULTRASAM_PATH}/models/ultrasam.py"
    )

# === 3️⃣ UltraSam 封装类 ===
class UltraSamWrapper:
    """
    UltraSam 模型推理封装
    - 自动加载权重 UltraSam.pth
    - 输入超声图像路径
    - 输出分割 mask（numpy 数组）
    """

    def __init__(self, device="cuda"):
        self.device = device
        ckpt_path = ULTRASAM_PATH / "UltraSam.pth"
        self.model = UltraSAM(weight_path=str(ckpt_path), device=self.device, input_size=(512,512))

    def preprocess(self, image_path: str):
        """图像预处理（根据 UltraSam 预期输入大小调整）"""
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"无法读取图像: {image_path}")
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, (512, 512))  # 统一大小，可按 UltraSam 要求调整
        img = img.astype(np.float32) / 255.0
        img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0)  # (1, 3, H, W)
        return img.to(self.device)

    def segment(self, image_path: str, save_path=None):
        mask = self.model.infer(image_path, save_path=save_path)
        return mask
